import pandas as pd
import numpy as np
from imblearn.ensemble import BalancedRandomForestClassifier
from sklearn.metrics import (classification_report, confusion_matrix, 
                            accuracy_score, precision_score, recall_score, 
                            f1_score, roc_auc_score, roc_curve, precision_recall_curve)
from sklearn.model_selection import GridSearchCV, cross_val_score
import matplotlib.pyplot as plt
import seaborn as sns
import warnings
warnings.filterwarnings('ignore')

# 设置中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

def load_processed_data():
    """加载预处理后的数据"""
    print("=== 加载数据 ===")
    
    try:
        # 加载训练数据和测试数据
        train_data = pd.read_excel('data_train_balanced.xlsx')
        test_data = pd.read_excel('data_test_validation.xlsx')
        
        print(f"训练集: {len(train_data)}行, {len(train_data.columns)}列")
        print(f"测试集: {len(test_data)}行, {len(test_data.columns)}列")
        
        # 分离特征和标签
        feature_columns = [col for col in train_data.columns if col != '是否异常']
        
        X_train = train_data[feature_columns]
        y_train = train_data['是否异常']
        X_test = test_data[feature_columns]
        y_test = test_data['是否异常']
        
        print(f"特征数量: {len(feature_columns)}")
        print(f"训练集分布 - 正常: {(y_train == 0).sum()}, 异常: {(y_train == 1).sum()}")
        print(f"测试集分布 - 正常: {(y_test == 0).sum()}, 异常: {(y_test == 1).sum()}")
        
        return X_train, X_test, y_train, y_test, feature_columns
        
    except Exception as e:
        print(f"数据加载失败: {e}")
        return None, None, None, None, None

def train_balanced_random_forest(X_train, y_train, random_state=42):
    """训练平衡随机森林分类器"""
    print("\n=== 训练平衡随机森林分类器 ===")
    
    # 基础模型
    balanced_rf = BalancedRandomForestClassifier(
        n_estimators=100,
        random_state=random_state,
        n_jobs=-1,
        class_weight=None,  # BalancedRandomForestClassifier内置处理不平衡
        bootstrap=True,
        oob_score=True
    )
    
    # 训练模型
    balanced_rf.fit(X_train, y_train)
    
    print("平衡随机森林训练完成")
    print(f"OOB得分: {balanced_rf.oob_score_:.4f}")
    
    return balanced_rf

def optimize_balanced_random_forest(X_train, y_train, random_state=42):
    """优化平衡随机森林超参数"""
    print("\n=== 优化平衡随机森林超参数 ===")
    
    # 参数网格 - 简化版本以提高运行速度
    param_grid = {
        'n_estimators': [100, 200],
        'max_depth': [None, 10, 20],
        'min_samples_split': [2, 5],
        'min_samples_leaf': [1, 2],
        'max_features': ['sqrt', 'log2']
    }
    
    # 基础模型
    balanced_rf = BalancedRandomForestClassifier(
        random_state=random_state,
        n_jobs=-1,
        bootstrap=True  # 确保bootstrap=True以支持oob_score
    )
    
    # 网格搜索
    grid_search = GridSearchCV(
        balanced_rf,
        param_grid,
        cv=5,
        scoring='f1',  # 使用F1得分作为评估指标
        n_jobs=-1,
        verbose=1
    )
    
    print("开始网格搜索...")
    grid_search.fit(X_train, y_train)
    
    print(f"最佳参数: {grid_search.best_params_}")
    print(f"最佳CV得分: {grid_search.best_score_:.4f}")
    
    return grid_search.best_estimator_

def evaluate_model_performance(model, X_train, X_test, y_train, y_test, model_name="平衡随机森林"):
    """评估模型性能"""
    print(f"\n=== 评估{model_name}性能 ===")
    
    # 预测
    y_train_pred = model.predict(X_train)
    y_test_pred = model.predict(X_test)
    
    # 预测概率
    y_train_proba = model.predict_proba(X_train)[:, 1]
    y_test_proba = model.predict_proba(X_test)[:, 1]
    
    # 计算各种指标
    results = {
        '数据集': ['训练集', '测试集'],
        '准确率': [
            accuracy_score(y_train, y_train_pred),
            accuracy_score(y_test, y_test_pred)
        ],
        '精确率': [
            precision_score(y_train, y_train_pred),
            precision_score(y_test, y_test_pred)
        ],
        '召回率': [
            recall_score(y_train, y_train_pred),
            recall_score(y_test, y_test_pred)
        ],
        'F1得分': [
            f1_score(y_train, y_train_pred),
            f1_score(y_test, y_test_pred)
        ],
        'AUC': [
            roc_auc_score(y_train, y_train_proba),
            roc_auc_score(y_test, y_test_proba)
        ]
    }
    
    results_df = pd.DataFrame(results)
    print(results_df.round(4))
    
    # 详细的测试集分类报告
    print(f"\n{model_name}测试集详细分类报告:")
    print(classification_report(y_test, y_test_pred, target_names=['正常', '异常']))
    
    # 混淆矩阵
    cm_train = confusion_matrix(y_train, y_train_pred)
    cm_test = confusion_matrix(y_test, y_test_pred)
    
    print(f"\n训练集混淆矩阵:")
    print(f"真阴性: {cm_train[0,0]}, 假阳性: {cm_train[0,1]}")
    print(f"假阴性: {cm_train[1,0]}, 真阳性: {cm_train[1,1]}")
    
    print(f"\n测试集混淆矩阵:")
    print(f"真阴性: {cm_test[0,0]}, 假阳性: {cm_test[0,1]}")
    print(f"假阴性: {cm_test[1,0]}, 真阳性: {cm_test[1,1]}")
    
    return results_df, y_test_pred, y_test_proba, cm_test

def analyze_feature_importance(model, feature_columns, top_k=15):
    """分析特征重要性"""
    print(f"\n=== 分析特征重要性（前{top_k}个） ===")
    
    # 获取特征重要性
    feature_importance = pd.DataFrame({
        '特征': feature_columns,
        '重要性': model.feature_importances_
    }).sort_values('重要性', ascending=False)
    
    # 显示前top_k个重要特征
    print(feature_importance.head(top_k))
    
    # 可视化特征重要性
    plt.figure(figsize=(10, 8))
    top_features = feature_importance.head(top_k)
    plt.barh(range(len(top_features)), top_features['重要性'])
    plt.yticks(range(len(top_features)), top_features['特征'])
    plt.xlabel('重要性')
    plt.title('平衡随机森林特征重要性分析')
    plt.gca().invert_yaxis()
    
    plt.tight_layout()
    plt.savefig('balanced_rf_feature_importance.png', dpi=300, bbox_inches='tight')
    plt.show()
    
    return feature_importance

def plot_roc_and_pr_curves(y_test, y_test_proba, model_name="平衡随机森林"):
    """绘制ROC曲线和PR曲线"""
    print(f"\n=== 生成{model_name}ROC曲线和PR曲线 ===")
    
    # 计算ROC曲线
    fpr, tpr, roc_thresholds = roc_curve(y_test, y_test_proba)
    roc_auc = roc_auc_score(y_test, y_test_proba)
    
    # 计算PR曲线
    precision, recall, pr_thresholds = precision_recall_curve(y_test, y_test_proba)
    
    # 创建子图
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
    
    # ROC曲线
    ax1.plot(fpr, tpr, 'b-', linewidth=2, label=f'{model_name} (AUC = {roc_auc:.4f})')
    ax1.plot([0, 1], [0, 1], 'r--', linewidth=2, label='随机分类器')
    ax1.set_xlim([0.0, 1.0])
    ax1.set_ylim([0.0, 1.05])
    ax1.set_xlabel('假阳性率')
    ax1.set_ylabel('真阳性率')
    ax1.set_title('ROC曲线')
    ax1.legend(loc="lower right")
    ax1.grid(True, alpha=0.3)
    
    # PR曲线
    ax2.plot(recall, precision, 'b-', linewidth=2, label=f'{model_name}')
    ax2.set_xlim([0.0, 1.0])
    ax2.set_ylim([0.0, 1.05])
    ax2.set_xlabel('召回率')
    ax2.set_ylabel('精确率')
    ax2.set_title('精确率-召回率曲线')
    ax2.legend()
    ax2.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('balanced_rf_roc_pr_curves.png', dpi=300, bbox_inches='tight')
    plt.show()

def plot_confusion_matrix(cm_test, model_name="平衡随机森林"):
    """绘制混淆矩阵热力图"""
    print(f"\n=== 生成{model_name}混淆矩阵 ===")
    
    plt.figure(figsize=(8, 6))
    sns.heatmap(cm_test, annot=True, fmt='d', cmap='Blues',
                xticklabels=['预测正常', '预测异常'],
                yticklabels=['实际正常', '实际异常'])
    plt.title(f'{model_name}测试集混淆矩阵')
    plt.ylabel('实际标签')
    plt.xlabel('预测标签')
    
    plt.tight_layout()
    plt.savefig('balanced_rf_confusion_matrix.png', dpi=300, bbox_inches='tight')
    plt.show()

def compare_sampling_strategies():
    """比较不同采样策略的效果"""
    print("\n=== 比较BalancedRandomForestClassifier的采样策略 ===")
    
    # 加载原始不平衡数据
    try:
        original_data = pd.read_excel('data_after_process_4.xlsx')
        feature_columns = [col for col in original_data.columns 
                         if col not in ['是否异常', '孕妇标识', 'Unnamed: 0']]
        
        X_original = original_data[feature_columns]
        y_original = original_data['是否异常']
        
        # 划分训练集和测试集
        from sklearn.model_selection import train_test_split
        X_train_orig, X_test_orig, y_train_orig, y_test_orig = train_test_split(
            X_original, y_original, test_size=0.2, random_state=42, stratify=y_original
        )
        
        print(f"原始数据分布 - 训练集: 正常{(y_train_orig==0).sum()}, 异常{(y_train_orig==1).sum()}")
        print(f"原始数据分布 - 测试集: 正常{(y_test_orig==0).sum()}, 异常{(y_test_orig==1).sum()}")
        
        # 训练BalancedRandomForestClassifier
        balanced_rf_orig = BalancedRandomForestClassifier(
            n_estimators=100,
            random_state=42,
            n_jobs=-1,
            oob_score=True
        )
        
        balanced_rf_orig.fit(X_train_orig, y_train_orig)
        y_pred_orig = balanced_rf_orig.predict(X_test_orig)
        
        print(f"\n使用原始不平衡数据的BalancedRandomForestClassifier性能:")
        print(f"准确率: {accuracy_score(y_test_orig, y_pred_orig):.4f}")
        print(f"精确率: {precision_score(y_test_orig, y_pred_orig):.4f}")
        print(f"召回率: {recall_score(y_test_orig, y_pred_orig):.4f}")
        print(f"F1得分: {f1_score(y_test_orig, y_pred_orig):.4f}")
        print(f"OOB得分: {balanced_rf_orig.oob_score_:.4f}")
        
        return balanced_rf_orig
        
    except Exception as e:
        print(f"比较采样策略时出错: {e}")
        return None

def generate_balanced_rf_report(results_df, cm_test, feature_importance, model_name="平衡随机森林"):
    """生成平衡随机森林详细报告"""
    print(f"\n=== 生成{model_name}分析报告 ===")
    
    # 计算额外指标
    tn, fp, fn, tp = cm_test.ravel()
    specificity = tn / (tn + fp) if (tn + fp) > 0 else 0
    sensitivity = tp / (tp + fn) if (tp + fn) > 0 else 0
    npv = tn / (tn + fn) if (tn + fn) > 0 else 0  # 阴性预测值
    ppv = tp / (tp + fp) if (tp + fp) > 0 else 0  # 阳性预测值（精确率）
    
    report = f"""
# 平衡随机森林分类器分析报告

## 模型概述
平衡随机森林分类器(BalancedRandomForestClassifier)是一种专门处理不平衡数据集的集成学习方法。
它通过在每棵决策树的训练过程中自动对多数类进行下采样，实现类别平衡。

## 模型性能指标

### 总体性能
- 训练集准确率: {results_df.loc[0, '准确率']:.4f}
- 测试集准确率: {results_df.loc[1, '准确率']:.4f}
- 训练集F1得分: {results_df.loc[0, 'F1得分']:.4f}
- 测试集F1得分: {results_df.loc[1, 'F1得分']:.4f}
- 测试集AUC: {results_df.loc[1, 'AUC']:.4f}

### 详细分类性能（测试集）
- **精确率**: {results_df.loc[1, '精确率']:.4f} - 预测为异常的样本中实际异常的比例
- **召回率**: {results_df.loc[1, '召回率']:.4f} - 实际异常样本中被正确识别的比例
- **特异性**: {specificity:.4f} - 实际正常样本中被正确识别的比例
- **阳性预测值(PPV)**: {ppv:.4f} - 与精确率相同
- **阴性预测值(NPV)**: {npv:.4f} - 预测为正常的样本中实际正常的比例

### 混淆矩阵分析（测试集）
- 真阴性(TN): {tn} - 正确识别为正常的样本数
- 假阳性(FP): {fp} - 误诊为异常的正常样本数
- 假阴性(FN): {fn} - 漏诊的异常样本数  
- 真阳性(TP): {tp} - 正确识别为异常的样本数

### 错误率分析
- 第一类错误(α): {fp/(tn+fp):.4f} - 将正常误判为异常的概率
- 第二类错误(β): {fn/(tp+fn):.4f} - 将异常误判为正常的概率（漏诊率）

## 特征重要性分析

### 最重要的5个特征
"""
    
    # 添加前5个重要特征
    for i, row in feature_importance.head(5).iterrows():
        report += f"- {row['特征']}: {row['重要性']:.4f}\n"
    
    report += f"""

## 平衡随机森林的优势

### 自动处理类别不平衡
- 每棵树训练时自动对多数类进行下采样
- 不需要额外的数据预处理步骤
- 保持了随机森林的所有优点

### 与传统随机森林的区别
- 传统随机森林在不平衡数据上倾向于预测多数类
- 平衡随机森林通过内置采样机制改善少数类的识别能力
- 减少了对外部采样技术(如SMOTE)的依赖

## 医学应用价值

### 适用场景
- 疾病筛查：识别高风险孕妇
- 临床决策支持：辅助医生判断
- 预警系统：早期发现异常情况

### 临床意义
- 召回率 {results_df.loc[1, '召回率']:.1%}: 能识别 {results_df.loc[1, '召回率']:.1%} 的异常病例
- 精确率 {results_df.loc[1, '精确率']:.1%}: 预测异常中有 {results_df.loc[1, '精确率']:.1%} 确实异常
- 漏诊率 {fn/(tp+fn):.1%}: 仅有 {fn/(tp+fn):.1%} 的异常病例被漏诊

## 模型部署建议

### 实际应用考虑
1. **阈值调整**: 可根据具体需求调整分类阈值
2. **定期更新**: 建议定期使用新数据重新训练模型
3. **专家验证**: 结合临床专家知识进行最终判断
4. **多模型融合**: 可考虑与其他模型结果进行融合

### 风险控制
- 建立模型监控机制
- 设置异常检测告警
- 保持人工审核机制

生成时间: {pd.Timestamp.now().strftime('%Y-%m-%d %H:%M:%S')}
"""
    
    # 保存报告
    with open('balanced_rf_analysis_report.txt', 'w', encoding='utf-8') as f:
        f.write(report)
    
    print("平衡随机森林分析报告已保存到: balanced_rf_analysis_report.txt")
    
    return report

def main():
    """主函数"""
    print("=== 平衡随机森林分类器分析 ===")
    
    # 1. 加载数据
    X_train, X_test, y_train, y_test, feature_columns = load_processed_data()
    if X_train is None:
        return
    
    # 2. 训练基础平衡随机森林
    basic_model = train_balanced_random_forest(X_train, y_train)
    
    # 3. 评估基础模型
    basic_results, y_pred_basic, y_proba_basic, cm_basic = evaluate_model_performance(
        basic_model, X_train, X_test, y_train, y_test, "基础平衡随机森林"
    )
    
    # 4. 优化模型超参数
    optimized_model = optimize_balanced_random_forest(X_train, y_train)
    
    # 5. 评估优化模型
    opt_results, y_pred_opt, y_proba_opt, cm_opt = evaluate_model_performance(
        optimized_model, X_train, X_test, y_train, y_test, "优化平衡随机森林"
    )
    
    # 6. 分析特征重要性
    feature_importance = analyze_feature_importance(optimized_model, feature_columns)
    
    # 7. 生成可视化
    plot_roc_and_pr_curves(y_test, y_proba_opt, "优化平衡随机森林")
    plot_confusion_matrix(cm_opt, "优化平衡随机森林")
    
    # 8. 比较采样策略
    original_model = compare_sampling_strategies()
    
    # 9. 生成详细报告
    report = generate_balanced_rf_report(opt_results, cm_opt, feature_importance)
    
    # 10. 输出总结
    print("\n" + "="*60)
    print("平衡随机森林分析完成！")
    print(f"优化后测试集F1得分: {opt_results.loc[1, 'F1得分']:.4f}")
    print(f"优化后测试集召回率: {opt_results.loc[1, '召回率']:.4f}")
    print(f"优化后测试集精确率: {opt_results.loc[1, '精确率']:.4f}")
    print("\n生成的文件:")
    print("📈 balanced_rf_feature_importance.png - 特征重要性图")
    print("📊 balanced_rf_roc_pr_curves.png - ROC和PR曲线")
    print("📊 balanced_rf_confusion_matrix.png - 混淆矩阵")
    print("📄 balanced_rf_analysis_report.txt - 详细分析报告")
    print("="*60)

if __name__ == "__main__":
    main()